This code is used to predict the team and player number from an image.
%run helpers.py
%run classification.py
%run transformations.py
%run model.py
%run prediction.py
import argparse
import os
import torch
import classification
import helpers
import transformations as tr
import model
import prediction
from torchvision import transforms
from torch.utils.data import DataLoader
def parse_cmdline_arguments():
"""Parse command line arguments"""
parser = argparse.ArgumentParser(description='Player and Team Recognition...')
parser.add_argument('--perform', type=str.lower,
choices=['train', 'predict', 'predict_and_color', 'train_and_predict'],
default='train_and_predict', help='Perform training or prediction or both')
parser.add_argument('-d', '--dir_path', type=str, default=(os.getcwd() + '\\data\\part1'),
help='Full path of dataset containing images for training the model')
parser.add_argument('-t', '--test_images_path', type=str, default=(os.getcwd() + '\\data\\test_images'),
help='Full path of dataset containing test images')
parser.add_argument('--csv_file_dir', type=str, default=(os.getcwd() + '\\data\\part2'),
help='CSV file directory having bbox for images')
parser.add_argument('-m', '--pretrained_model', type=str.lower,
choices=['resnet18', 'resnet34', 'resnet50', 'densenet161', 'inceptionresnetv2'],
default='resnet18', help='Pretrained base model')
parser.add_argument('-e', '--epochs', type=int, default=20, help='No. of epochs')
parser.add_argument('--optimizer', type=str.lower, choices=['sgd', 'adam'],
default='sgd', help='Optimizer function')
parser.add_argument('-bs', '--batch_size', type=int, default=32, help='Batch size')
parser.add_argument('-lr', '--learning_rate', type=float, default=0.0001, help='Learning rate')
parser.add_argument('--show_plot', type=bool, choices=[True, False], default=True,
help='Display/Show corresponding plot')
return parser.parse_args(args=[])
def main():
args = parse_cmdline_arguments()
# Check if the dataset directory exists and import the information about the players
helpers.check_path_exists(args.dir_path, 'train images directory')
df_player_info, teams_dic, players_dic = helpers.generate_player_info(args.dir_path, check_invalid_images_flag=False)
# Split the data and apply the transformations (resize, normalize, etc)
data = helpers.split_data(df_player_info, test_size=0.2)
train_transform = transforms.Compose([tr.Resize(224, 224), tr.ToTensor(), tr.Normalization()])
test_transform = transforms.Compose([tr.Resize(224, 224), tr.ToTensor(), tr.Normalization()])
data['train'] = classification.Classification(player_info=data['train'], transform=train_transform)
data['valid'] = classification.Classification(player_info=data['valid'], transform=train_transform)
# Generate the dataloaders for training and validation set
train_dataloader = DataLoader(dataset=data['train'], batch_size=args.batch_size, shuffle=True, num_workers=0)
valid_dataloader = DataLoader(dataset=data['valid'], batch_size=args.batch_size, shuffle=False, num_workers=0)
dataloaders = {phase:train_dataloader if phase == 'train' else valid_dataloader for phase in ['train', 'valid']}
# Define the model and move to the device available
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
model_CNN = model.Model(True, args.pretrained_model, len(teams_dic), len(players_dic), data, args)
# Define the Optimizer and Loss function
criterions, optimizer = model_CNN.compilation_parameters(model_CNN, args)
# Perform the model training
if args.perform == 'train' or args.perform == 'train_and_predict':
model_CNN = model_CNN.train_model(model_CNN, dataloaders, criterions, optimizer, device, n_epochs=args.epochs,
show_plot = args.show_plot)
torch.save(model_CNN.state_dict(), 'model_CNN.ckpt')
# Perform prediction on the images present in the specified directory
if args.perform == 'predict' or args.perform == 'train_and_predict':
model_CNN.load_state_dict(torch.load('model_CNN.ckpt'))
predict = prediction.Prediction(model_CNN, test_transform, teams_dic, players_dic)
helpers.check_path_exists(args.test_images_path, 'test images directory')
predict.predict_images_in_directory(args.test_images_path, check_invalid_images_flag=False)
# Perform prediction and color the section of images as per the prediction outcome
if args.perform == 'predict_and_color' or args.perform == 'train_and_predict':
model_CNN.load_state_dict(torch.load('model_CNN.ckpt'))
predict = prediction.Prediction(model_CNN, test_transform, teams_dic, players_dic)
predict.predict_and_color_section_of_images(args.csv_file_dir, show_plot=args.show_plot)
if __name__ == '__main__':
main()
Epoch: 1 Training Acc: 0.0175 Validation Acc: 0.0534 Epoch: 2 Training Acc: 0.0534 Validation Acc: 0.0850 Epoch: 3 Training Acc: 0.0577 Validation Acc: 0.0882 Epoch: 4 Training Acc: 0.0560 Validation Acc: 0.0894 Epoch: 5 Training Acc: 0.0522 Validation Acc: 0.0963 Epoch: 6 Training Acc: 0.0477 Validation Acc: 0.1094 Epoch: 7 Training Acc: 0.0444 Validation Acc: 0.1097 Epoch: 8 Training Acc: 0.0432 Validation Acc: 0.1149 Epoch: 9 Training Acc: 0.0407 Validation Acc: 0.1294 Epoch: 10 Training Acc: 0.0406 Validation Acc: 0.1498 Epoch: 11 Training Acc: 0.0419 Validation Acc: 0.1620 Epoch: 12 Training Acc: 0.0441 Validation Acc: 0.1789 Epoch: 13 Training Acc: 0.0445 Validation Acc: 0.1907 Epoch: 14 Training Acc: 0.0457 Validation Acc: 0.1982 Epoch: 15 Training Acc: 0.0474 Validation Acc: 0.2016 Epoch: 16 Training Acc: 0.0471 Validation Acc: 0.2168 Epoch: 17 Training Acc: 0.0492 Validation Acc: 0.2196 Epoch: 18 Training Acc: 0.0499 Validation Acc: 0.2242 Epoch: 19 Training Acc: 0.0491 Validation Acc: 0.2223 Epoch: 20 Training Acc: 0.0501 Validation Acc: 0.2198
Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\bristol_person_14.jpg Predictions - Team: middlesbrough, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\bristol_person_2.jpg Predictions - Team: wigan, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\bristol_person_25.jpg Predictions - Team: wigan, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\bristol_person_5.jpg Predictions - Team: wigan, Player_no: person_5 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\middlesbrough_person_11.jpg Predictions - Team: middlesbrough, Player_no: person_6 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\middlesbrough_person_27.jpg Predictions - Team: middlesbrough, Player_no: person_6 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\middlesbrough_person_5.jpg Predictions - Team: middlesbrough, Player_no: person_6 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\nottingham_forrest_person_3.jpg Predictions - Team: bristol, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\nottingham_forrest_person_8.jpg Predictions - Team: bristol, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\spal_team_a_person_10.jpg Predictions - Team: spal_team_a, Player_no: person_5 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\spal_team_a_person_19.jpg Predictions - Team: spal_team_a, Player_no: person_19 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\spal_team_a_person_37.jpg Predictions - Team: wigan, Player_no: person_19 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\spal_team_a_person_77.jpg Predictions - Team: spal_team_a, Player_no: person_19 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\spal_team_b_person_23.jpg Predictions - Team: wigan, Player_no: person_5 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\wigan_person_11.jpg Predictions - Team: bristol, Player_no: person_8 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\wigan_person_19.jpg Predictions - Team: bristol, Player_no: person_5 Test image path: d:\OneDrive\Desktop\player_and_team_recognition\data\test_images\wigan_person_22.jpg Predictions - Team: bristol, Player_no: person_8 CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\11.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\12.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\13.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\14.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\15.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\16.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\17.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\18.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\19.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\20.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\21.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\22.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\23.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\24.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\25.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\26.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\27.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\28.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\29.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\30.csv
CSV File: d:\OneDrive\Desktop\player_and_team_recognition\data\part2\4.csv
Few references/weblinks:
Some of these tasks has been with Tensorflow, but can be easily done using PyTorch framework.